"""
python p2_solution_deap.py
"""
import json
import os
import random
from deap import base, creator, tools, algorithms

# Define paths
warehouse_path = '../fujian/fujian3/origin_data/warehouse.json'
inventory_path = '../fujian/fujian3/data_from_p1/all_average_inventory.json'
sales_path = '../fujian/fujian3/data_from_p1/all_average_sales.json'
output_path = '../fujian/fujian3/deap/'
os.makedirs(output_path, exist_ok=True)

# Load data
with open(warehouse_path, 'r') as f:
    warehouses = json.load(f)
with open(inventory_path, 'r') as f:
    inventory_data = json.load(f)
with open(sales_path, 'r') as f:
    sales_data = json.load(f)

# Process data into dictionaries for quick access
warehouse_dict = {w['warehouse_id']: w for w in warehouses}
inventory_dict = {i['category_id']: i['average_inventory'] for i in inventory_data}
sales_dict = {s['category_id']: s['average_sales'] for s in sales_data}

categories = list(inventory_dict.keys())
num_categories = len(categories)
warehouses_list = list(warehouse_dict.keys())
num_warehouses = len(warehouses_list)

# Create a fitness function to minimize cost and check constraints
def evaluate(individual):
    total_cost = 0
    warehouse_usage = {w: {'inventory': 0, 'sales': 0} for w in warehouses_list}

    # Calculate usage and cost
    for category_idx, warehouse_id in enumerate(individual):
        category_id = categories[category_idx]
        inventory = inventory_dict[category_id]
        sales = sales_dict[category_id]
        warehouse = warehouse_dict[warehouse_id]

        # Update warehouse utilization
        warehouse_usage[warehouse_id]['inventory'] += inventory
        warehouse_usage[warehouse_id]['sales'] += sales

        # Add cost if this warehouse is used
        if inventory > 0:
            total_cost += warehouse['daily_cost']

    # Check constraints and penalize if constraints are violated
    penalty = 0
    for warehouse_id, usage in warehouse_usage.items():
        warehouse = warehouse_dict[warehouse_id]
        if usage['inventory'] > warehouse['max_inventory']:
            penalty += 1000 * (usage['inventory'] - warehouse['max_inventory'])
        if usage['sales'] > warehouse['max_sales']:
            penalty += 1000 * (usage['sales'] - warehouse['max_sales'])

    return total_cost + penalty,

# DEAP setup
creator.create("FitnessMin", base.Fitness, weights=(-1.0,))  # Minimize cost
creator.create("Individual", list, fitness=creator.FitnessMin)

toolbox = base.Toolbox()
toolbox.register("indices", random.choices, warehouses_list, k=num_categories)
toolbox.register("individual", tools.initIterate, creator.Individual, toolbox.indices)
toolbox.register("population", tools.initRepeat, list, toolbox.individual)
toolbox.register("evaluate", evaluate)
toolbox.register("mate", tools.cxTwoPoint)
toolbox.register("mutate", tools.mutFlipBit, indpb=0.05)
toolbox.register("select", tools.selTournament, tournsize=3)

# Run the genetic algorithm
population = toolbox.population(n=100)
ngen = 100
for gen in range(ngen):
    # Select the next generation individuals
    offspring = toolbox.select(population, len(population))
    offspring = list(map(toolbox.clone, offspring))

    # Apply crossover and mutation
    for child1, child2 in zip(offspring[::2], offspring[1::2]):
        if random.random() < 0.5:
            toolbox.mate(child1, child2)
            del child1.fitness.values
            del child2.fitness.values

    for mutant in offspring:
        if random.random() < 0.2:
            toolbox.mutate(mutant)
            del mutant.fitness.values

    # Evaluate individuals with an invalid fitness
    invalid_ind = [ind for ind in offspring if not ind.fitness.valid]
    fitnesses = map(toolbox.evaluate, invalid_ind)
    for ind, fit in zip(invalid_ind, fitnesses):
        ind.fitness.values = fit

    # Replace the old population by the offspring
    population[:] = offspring

# Get the best individual
fits = [ind.fitness.values[0] for ind in population]
best_idx = fits.index(min(fits))
best_individual = population[best_idx]

# Prepare output based on best_individual
warehouse_utilization = {w: {'inventory': 0, 'sales': 0} for w in warehouses_list}
assignment = []

for category_idx, warehouse_id in enumerate(best_individual):
    category_id = categories[category_idx]
    inventory = inventory_dict[category_id]
    sales = sales_dict[category_id]
    
    # Update utilization
    warehouse_utilization[warehouse_id]['inventory'] += inventory
    warehouse_utilization[warehouse_id]['sales'] += sales
    
    # Store assignment
    assignment.append({'category_id': category_id, 'warehouse_id': warehouse_id})

# Save assignment
with open(os.path.join(output_path, 'assignment.json'), 'w') as f:
    json.dump(assignment, f)

# Calculate total cost
total_cost = evaluate(best_individual)[0]

# Save total cost
with open(os.path.join(output_path, 'all_cost.txt'), 'w') as f:
    f.write(str(total_cost))

# Calculate and save utilization rates
utilization_rates_inventory = []
utilization_rates_sales = []

for warehouse_id, usage in warehouse_utilization.items():
    max_inventory = warehouse_dict[warehouse_id]['max_inventory']
    max_sales = warehouse_dict[warehouse_id]['max_sales']
    
    utilization_rates_inventory.append({
        'warehouse_id': warehouse_id,
        'utilization_rate_of_inventory': usage['inventory'] / max_inventory
    })
    utilization_rates_sales.append({
        'warehouse_id': warehouse_id,
        'utilization_rate_of_sales': usage['sales'] / max_sales
    })

# Save utilization rates to JSON files
with open(os.path.join(output_path, 'all_utilization_rate_inventory.json'), 'w') as f:
    json.dump(utilization_rates_inventory, f)
with open(os.path.join(output_path, 'all_utilization_rate_sales.json'), 'w') as f:
    json.dump(utilization_rates_sales, f)

print("Optimization Complete. Best total cost:", total_cost)
